#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import time
import psutil
import parted
import struct
import subprocess


class Util:

    @staticmethod
    def formatSize(value):
        # value is in bytes
        if value > 1024 * 1024 * 1024 * 1024:
            return "%.1fTiB" % (value / 1024 / 1024 / 1024 / 1024)
        elif value > 1024 * 1024 * 1024:
            return "%.1fGiB" % (value / 1024 / 1024 / 1024)
        elif value > 1024 * 1024:
            return "%.1fMiB" % (value / 1024 / 1024)
        elif value > 1024:
            return "%.1fKiB" % (value / 1024)
        else:
            assert False

    @staticmethod
    def wipeHarddisk(devpath):
        # write data to harddisk
        fd = os.open(devpath, os.O_WRONLY | os.O_EXCL)
        try:
            for i in range(0, 1024):
                os.write(fd, bytearray(4096))
        finally:
            os.close(fd)

        # wait for /dev refresh
        while PartiUtil.inScope(devpath) and PartiUtil.diskHasParti(devpath, 1):
            print("FIXME: %s, still has partition" % (devpath))
            time.sleep(1)

    @staticmethod
    def getBlkDevFsUuid(devPath):
        try:
            out = subprocess.check_output(["blkid", "-o", "export", devPath], universal_newlines=True)
            m = re.search("^UUID=(\\S+)$", out, re.M)
            if m is not None:
                return m.group(1)
        except subprocess.CalledProcessError as e:
            if e.returncode != 2:
                raise
        return ""

    @staticmethod
    def getBlkDevFsTypeAndFsUuidAndLabel(devPath):
        ret = ["", "", ""]

        try:
            out = subprocess.check_output(["blkid", "-o", "export", devPath], universal_newlines=True)

            m = re.search("^TYPE=(\\S+)$", out, re.M)
            if m is not None:
                ret[0] = m.group(1).lower()
            else:
                return ret

            m = re.search("^UUID=(\\S+)$", out, re.M)
            ret[1] = m.group(1)                                 # FS-UUID must exist if we have fstype

            m = re.search("^LABEL=(\\S+)$", out, re.M)
            if m is not None:
                ret[2] = m.group(1)                             # LABEL is optional

        except subprocess.CalledProcessError as e:
            if e.returncode != 2:
                raise

        return ret

    @staticmethod
    def getBlkDevSize(devPath):
        out = subprocess.check_output(["blockdev", "--getsz", devPath], universal_newlines=True)
        return int(out) * 512        # unit is byte

    @staticmethod
    def getBlkDevCapacity(devPath):
        ret = subprocess.check_output(["df", "-BM", devPath], universal_newlines=True)
        m = re.search("%s +(\\d+)M +(\\d+)M +\\d+M", ret, re.M)
        total = int(m.group(1))        # unit: MB
        used = int(m.group(2))         # unit: MB
        return (total * 1024 * 1024, used * 1024 * 1024)

    @staticmethod
    def isHarddiskBusy(devpath):
        try:
            fd = os.open(devpath, os.O_WRONLY | os.O_EXCL)
            os.close(fd)
        except OSError as e:
            if e.errno == 16:
                return True
            else:
                raise
        return False

    @staticmethod
    def initializeDisk(devPath, partitionTableType, partitionInfoList):
        assert PartiUtil.inScope(devPath)
        assert partitionTableType in ["mbr", "msdos"]
        assert len(partitionInfoList) >= 1

        def _getFreeRegion(disk):
            region = None
            for r in disk.getFreeSpaceRegions():
                if r.length <= disk.device.optimumAlignment.grainSize:
                    continue                                                # ignore alignment gaps
                if region is not None:
                    assert False                                            # there should be only one free region
                region = r
            return region

        def _addPartition(disk, pType, pStart, pEnd):
            region = parted.Geometry(device=disk.device, start=pStart, end=pEnd)
            if pType == "":
                partition = parted.Partition(disk=disk, type=parted.PARTITION_NORMAL, geometry=region)
            elif pType in ["vfat", "exfat"]:
                partition = parted.Partition(disk=disk,
                                             type=parted.PARTITION_NORMAL,
                                             fs=parted.FileSystem(type="fat32", geometry=region),
                                             geometry=region)
            else:
                assert False
            disk.addPartition(partition=partition,
                              constraint=disk.device.optimalAlignedConstraint)

        def _erasePartitionSignature(devPath, pStart, pEnd):
            # fixme: this implementation is very limited
            with open(devPath, "wb") as f:
                f.seek(pStart * 512)
                if pEnd - pStart + 1 < 32:
                    f.write(bytearray((pEnd - pStart + 1) * 512))
                else:
                    f.write(bytearray(32 * 512))

        # partitionInfoList => preList & postList
        preList = None
        postList = None
        for i in range(0, len(partitionInfoList)):
            pSize, pType = partitionInfoList[i]
            if pSize == "*":
                assert preList is None
                preList = partitionInfoList[:i]
                postList = partitionInfoList[i:]
        if preList is None:
            preList = partitionInfoList
            postList = []

        # sucks that libparted does not support open device exclusively
        assert not Util.isHarddiskBusy(devPath)

        # delete all partitions
        # 1. we must do it manually because we need a clean /dev directory to do checks later
        # 2. it seems parted does not clean the boot-code, this function should clean it
        Util.wipeHarddisk(devPath)

        # create new disk object
        disk = parted.freshDisk(parted.getDevice(devPath), partitionTableType)

        # process preList
        for pSize, pType in preList:
            region = _getFreeRegion(disk)
            constraint = parted.Constraint(maxGeom=region).intersect(disk.device.optimalAlignedConstraint)
            pStart = constraint.startAlign.alignUp(region, region.start)
            pEnd = constraint.endAlign.alignDown(region, region.end)

            m = re.fullmatch("([0-9]+)(MiB|GiB|TiB)", pSize)
            assert m is not None
            sectorNum = parted.sizeToSectors(int(m.group(1)), m.group(2), disk.device.sectorSize)
            if pEnd < pStart + sectorNum - 1:
                raise Exception("not enough space")

            _addPartition(disk, pType, pStart, pStart + sectorNum - 1)
            _erasePartitionSignature(devPath, pStart, pEnd)

        # process postList
        for pSize, pType in postList:
            region = _getFreeRegion(disk)
            constraint = parted.Constraint(maxGeom=region).intersect(disk.device.optimalAlignedConstraint)
            pStart = constraint.startAlign.alignUp(region, region.start)
            pEnd = constraint.endAlign.alignDown(region, region.end)

            if pSize == "*":
                _addPartition(disk, pType, pStart, pEnd)
                _erasePartitionSignature(devPath, pStart, pEnd)
            else:
                assert False

        # write to disk, notify kernel (using BLKRRPART ioctl), block until kernel picks up this change
        disk.commit()

        # wait partition device nodes appear in /dev
        # there's still a time gap between kernel and /dev refresh, maybe because udevd?
        for i in range(0, len(partitionInfoList)):
            while not PartiUtil.diskHasParti(devPath, i + 1):
                print("FIXME: partition %d of %s does not exist" % (i + 1, devPath))
                time.sleep(1)

    @staticmethod
    def isBufferAllZero(buf):
        for b in buf:
            if b != 0:
                return False
        return True

    @staticmethod
    def newGuid(guidStr):
        assert len(guidStr) == 36
        assert guidStr[8] == "-" and guidStr[13] == "-" and guidStr[18] == "-" and guidStr[23] == "-"

        # struct gpt_guid {
        #     uint32_t   time_low;
        #     uint16_t   time_mid;
        #     uint16_t   time_hi_and_version;
        #     uint8_t    clock_seq_hi;
        #     uint8_t    clock_seq_low;
        #     uint8_t    node[6];
        # };
        gptGuidFmt = "IHHBB6s"
        assert struct.calcsize(gptGuidFmt) == 16

        guidStr = guidStr.replace("-", "")

        # really obscure behavior of python3
        # see http://stackoverflow.com/questions/1463306/how-does-exec-work-with-locals
        ldict = {}
        exec("n1 = 0x" + guidStr[0:8], globals(), ldict)
        exec("n2 = 0x" + guidStr[8:12], globals(), ldict)
        exec("n3 = 0x" + guidStr[12:16], globals(), ldict)
        exec("n4 = 0x" + guidStr[16:18], globals(), ldict)
        exec("n5 = 0x" + guidStr[18:20], globals(), ldict)
        exec("n6 = bytearray()", globals(), ldict)
        for i in range(0, 6):
            exec("n6.append(0x" + guidStr[20 + i * 2:20 + (i + 1) * 2] + ")", globals(), ldict)

        return struct.pack(gptGuidFmt, ldict["n1"], ldict["n2"], ldict["n3"], ldict["n4"], ldict["n5"], ldict["n6"])


class PartiUtil:

    @staticmethod
    def inScope(devPath):
        if re.fullmatch("/dev/sd[a-z][0-9]*", devPath) is not None:
            return True
        if re.fullmatch("/dev/xvd[a-z][0-9]*", devPath) is not None:
            return True
        if re.fullmatch("/dev/vd[a-z][0-9]*", devPath) is not None:
            return True
        if re.fullmatch("/dev/nvme[0-9]+n[0-9]+p[0-9]*", devPath) is not None:
            return True
        return False

    @staticmethod
    def isDiskOrParti(devPath):
        if re.fullmatch("/dev/sd[a-z]", devPath) is not None:
            return True
        if re.fullmatch("(/dev/sd[a-z])([0-9]+)", devPath) is not None:
            return False
        if re.fullmatch("/dev/xvd[a-z]", devPath) is not None:
            return True
        if re.fullmatch("(/dev/xvd[a-z])([0-9]+)", devPath) is not None:
            return False
        if re.fullmatch("/dev/vd[a-z]", devPath) is not None:
            return True
        if re.fullmatch("(/dev/vd[a-z])([0-9]+)", devPath) is not None:
            return False
        if re.fullmatch("/dev/nvme[0-9]+n[0-9]+", devPath) is not None:
            return True
        if re.fullmatch("(/dev/nvme[0-9]+n[0-9]+)p([0-9]+)", devPath) is not None:
            return False
        assert False

    @staticmethod
    def partiToDiskAndPartiId(partitionDevPath):
        m = re.fullmatch("(/dev/sd[a-z])([0-9]+)", partitionDevPath)
        if m is not None:
            return (m.group(1), int(m.group(2)))
        m = re.fullmatch("(/dev/xvd[a-z])([0-9]+)", partitionDevPath)
        if m is not None:
            return (m.group(1), int(m.group(2)))
        m = re.fullmatch("(/dev/vd[a-z])([0-9]+)", partitionDevPath)
        if m is not None:
            return (m.group(1), int(m.group(2)))
        m = re.fullmatch("(/dev/nvme[0-9]+n[0-9]+)p([0-9]+)", partitionDevPath)
        if m is not None:
            return (m.group(1), int(m.group(2)))
        assert False

    @staticmethod
    def partiToDisk(partitionDevPath):
        return PartiUtil.partiToDiskAndPartiId(partitionDevPath)[0]

    @staticmethod
    def diskToParti(diskDevPath, partitionId):
        m = re.fullmatch("/dev/sd[a-z]", diskDevPath)
        if m is not None:
            return diskDevPath + str(partitionId)
        m = re.fullmatch("/dev/xvd[a-z]", diskDevPath)
        if m is not None:
            return diskDevPath + str(partitionId)
        m = re.fullmatch("/dev/vd[a-z]", diskDevPath)
        if m is not None:
            return diskDevPath + str(partitionId)
        m = re.fullmatch("/dev/nvme[0-9]+n[0-9]+", diskDevPath)
        if m is not None:
            return diskDevPath + "p" + str(partitionId)
        assert False

    @staticmethod
    def diskHasParti(diskDevPath, partitionId):
        partiDevPath = PartiUtil.diskToParti(diskDevPath, partitionId)
        return os.path.exists(partiDevPath)

    @staticmethod
    def diskHasMoreParti(diskDevPath, partitionId):
        for i in range(partitionId + 1, partitionId + 10):
            if os.path.exists(PartiUtil.diskToParti(diskDevPath, i)):
                return True
        return False

    @staticmethod
    def partiExists(partitionDevPath):
        return os.path.exists(partitionDevPath)


class PhysicalDiskMounts:

    """This class is a better psutil.disk_partitions()"""

    class Entry:

        def __init__(self, p):
            self._p = p

        @property
        def device(self):
            return self._p.device

        @property
        def mountpoint(self):
            return self._p.mountpoint

        @property
        def fstype(self):
            return self._p.fstype

        @property
        def opts(self):
            return self._p.opts

        @property
        def mnt_opt_list(self):
            return self._p.opts.split(",")

        def __repr__(self):
            return "<%s %r>" % (self.__class__.__name__, self.__dict__)

    class NotFoundError(Exception):
        pass

    @classmethod
    def get_entries(cls):
        return [cls.Entry(p) for p in psutil.disk_partitions()]

    @classmethod
    def get_single_entry(cls, device, mountpoint):
        ret = None
        for p in psutil.disk_partitions():
            if p.mountpoint == mountpoint:
                assert p.device == device
                assert ret is None
                ret = cls.Entry(p)
        return ret
